package io.kiki.sba.registry.server.clientInterface.push;


import io.kiki.sba.registry.api.Callback;
import io.kiki.sba.registry.api.Channel;
import io.kiki.sba.registry.api.ChannelOverflowException;
import io.kiki.sba.registry.api.exchange.RequestChannelClosedException;
import io.kiki.sba.registry.common.model.DataCenterPushInfo;
import io.kiki.sba.registry.common.model.Tuple;
import io.kiki.sba.registry.common.model.store.DataCenterToSubDatumMap;
import io.kiki.sba.registry.common.model.store.PushData;
import io.kiki.sba.registry.common.model.store.Subscriber;
import io.kiki.sba.registry.server.clientInterface.bootstrap.Config;
import io.kiki.sba.registry.server.clientInterface.circuit.breaker.CircuitBreakerService;
import io.kiki.sba.registry.server.clientInterface.node.service.ClientNodeService;
import io.kiki.sba.registry.server.shared.util.DatumUtils;
import io.kiki.sba.registry.task.KeyedThreadPoolExecutor;
import io.kiki.sba.registry.task.MetricsableThreadPoolExecutor;
import io.kiki.sba.registry.task.RejectedDiscardHandler;
import io.kiki.sba.registry.trace.TraceId;
import io.kiki.sba.registry.util.*;
import io.kiki.stack.netty.exception.InvokeTimeoutException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;

import javax.annotation.PostConstruct;
import java.net.InetSocketAddress;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import static io.kiki.sba.registry.server.clientInterface.push.PushMetrics.Push.*;

public class PushProcessor {
    private static final Logger logger = LoggerFactory.getLogger(PushProcessor.class);
    final Map<PushTask.PushingTaskKey, PushRecord> pushingPushingTaskKeyToRecordMap = new ConcurrentHashMap<>(1024 * 16);
    final Cleaner cleaner = new Cleaner();
    final RejectedDiscardHandler rejectedDiscardHandler = new RejectedDiscardHandler();
    private final ThreadPoolExecutor pushCallbackExecutor = MetricsableThreadPoolExecutor.newExecutor("PushCallback", OsUtils.getCpuCount() * 3, 40000, rejectedDiscardHandler);
    @Autowired
    protected Config config;
    @Autowired
    protected PushSwitchService pushSwitchService;
    @Autowired
    protected PushDataGenerator pushDataGenerator;
    @Autowired
    protected ClientNodeService clientNodeService;
    ;
    @Autowired
    protected CircuitBreakerService circuitBreakerService;
    PushTaskWorkers pushTaskWorkers;
    private KeyedThreadPoolExecutor pushExecutor;
    private int pushDataTaskDebouncingMillis = 500;
    private PushEfficiencyImproveConfig pushEfficiencyImproveConfig;

    @PostConstruct
    public void init() {
        pushExecutor = new KeyedThreadPoolExecutor("PushExecutor", config.getPushTaskExecutorPoolSize(), config.getPushTaskExecutorQueueSize());
        intPushTaskWorkers();
        ConcurrentUtils.createDaemonThread("PushCleaner", cleaner).start();
    }

    public void setPushTaskDelayTime(PushEfficiencyImproveConfig pushEfficiencyImproveConfig) {
        this.pushTaskWorkers.setPushTaskWorkWaitingMillis(pushEfficiencyImproveConfig.getPushTaskWaitingMillis());
        this.pushDataTaskDebouncingMillis = pushEfficiencyImproveConfig.getPushTaskDebouncingMillis();
        this.pushEfficiencyImproveConfig = pushEfficiencyImproveConfig;
    }

    void intPushTaskWorkers() {
        if (this.pushTaskWorkers == null) {
            this.pushTaskWorkers = new PushTaskWorkers(config.getPushTaskBufferBucketSize());
        }
    }

    protected List<PushTask> createPushTask(Cause cause, InetSocketAddress inetSocketAddress, Map<String, Subscriber> registerIdToSubscriberMap, DataCenterToSubDatumMap dataCenterToSubDatumMap) {
        PushTask pushTask = new PushTaskImpl(cause, inetSocketAddress, registerIdToSubscriberMap, dataCenterToSubDatumMap);
        // set expireTimestamp, wait to merge to debouncing
        if (null != pushEfficiencyImproveConfig) {
            pushTask.expireAfter(pushEfficiencyImproveConfig.fetchSbfAppPushTaskDebouncingMillis(pushTask.subscriber.getAppName()));
        } else {
            pushTask.expireAfter(pushDataTaskDebouncingMillis);
        }
        return Collections.singletonList(pushTask);
    }

    void firePush(Cause cause, InetSocketAddress inetSocketAddress, Map<String, Subscriber> registerIdToSubscriberMap, DataCenterToSubDatumMap dataCenterToSubDatumMap) {
        if (!pushSwitchService.canIpPushMulti(inetSocketAddress.getAddress().getHostAddress(), dataCenterToSubDatumMap.dataCenterSet())) {
            return;
        }
        // most of the time, element size is 1, SingleMap to save the memory
        registerIdToSubscriberMap = CollectionUtils.toSingletonMap(registerIdToSubscriberMap);
        List<PushTask> pushTaskList = createPushTask(cause, inetSocketAddress, registerIdToSubscriberMap, dataCenterToSubDatumMap);
        for (PushTask pushTask : pushTaskList) {
            pushTaskWorkers.add(pushTask);
        }
    }

    private int getPushingMaxSpanMillis() {
        return config.getClientNodeExchangeTimeoutMillis() * 3;
    }

    int cleanPushingTaskRunTooLong() {
        final long now = System.currentTimeMillis();
        final int maxSpanMillis = getPushingMaxSpanMillis();
        int count = 0;
        for (Map.Entry<PushTask.PushingTaskKey, PushRecord> entry : pushingPushingTaskKeyToRecordMap.entrySet()) {
            final PushTask.PushingTaskKey pushingTaskKey = entry.getKey();
            final PushRecord pushRecord = entry.getValue();
            long span = cleanPushingTaskIfRunTooLong(now, pushingTaskKey, pushRecord, maxSpanMillis);
            if (span > 0) {
                count++;
                logger.warn("[pushTooLong]{},span={},{}", pushRecord.traceId, span, pushingTaskKey);
            }
        }
        return count;
    }

    long cleanPushingTaskIfRunTooLong(long now, PushTask.PushingTaskKey pushingTaskKey, PushRecord pushRecord, int maxSpanMillis) {
        final long span = now - pushRecord.trace.getPushStartTimestamp();
        if (span > maxSpanMillis) {
            // this happens when the callbackExecutor is too busy and the callback task is discarded
            // force to remove the prev task
            final boolean cleaned = pushingPushingTaskKeyToRecordMap.remove(pushingTaskKey, pushRecord);
            if (cleaned) {
                for (DataCenterPushInfo dataCenterPushInfo : pushRecord.stringDataCenterPushInfoMap.values()) {
                    dataCenterPushInfo.setPushVersion(0L);
                }
                pushRecord.trace.finishPush(Trace.PushStatus.Busy, pushRecord.traceId, pushRecord.stringDataCenterPushInfoMap, pushRecord.retryCount);
            }
            return span;
        }
        return 0;
    }

    boolean interest(PushTask pushTask) {
        Map<String, Long> dataCenterToSubDatumVersionMap = pushTask.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap();
        if (pushTask.registerIdToSubscriberMap.size() == 1) {
            return pushTask.subscriber.checkVersion(dataCenterToSubDatumVersionMap);
        }
        for (Subscriber subscriber : pushTask.registerIdToSubscriberMap.values()) {
            if (subscriber.checkVersion(dataCenterToSubDatumVersionMap)) {
                return true;
            }
        }
        return false;
    }

    boolean checkPushRunning(PushTask task) {
        final PushTask.PushingTaskKey pushingTaskKey = task.pushingTaskKey;
        // check the pushing task
        final PushRecord prev = pushingPushingTaskKeyToRecordMap.get(pushingTaskKey);
        if (prev == null) {
            return true;
        }
        final long now = System.currentTimeMillis();
        final int maxSpanMillis = getPushingMaxSpanMillis();
        final long span = cleanPushingTaskIfRunTooLong(now, pushingTaskKey, prev, maxSpanMillis);
        if (span > 0) {
            logger.warn("[pushTooLong]{},span={},{},now={}", prev.traceId, span, task.pushingTaskKey, task.traceId);
            return true;
        }
        // task after the prev, but prev.pushClient not callback, retry
        retry(task, RetryReason.Waiting);
        return false;
    }

    boolean causeContinue(PushTask pushTask) {
        switch (pushTask.trace.cause.type) {
        case reg:
            return !pushTask.hasPushed();
        case empty:
            return pushTask.subscriber.needPushEmpty(pushTask.dataCenterToSubDatumMap.dataCenterSet());
        default:
            return interest(pushTask);
        }
    }

    boolean checkSkipPushEmptyAndUpdateVersion(PushTask task) {
        int skipCount = 0;
        long now = System.currentTimeMillis();
        Collection<Subscriber> subs = task.registerIdToSubscriberMap.values();
        for (Subscriber subscriber : subs) {
            boolean canSkip = subscriber.checkSkipPushEmpty(task.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), task.getDataCenterPushCount());
            if (canSkip && subscriber.getRegisterTimestamp() < now - config.getSkipPushEmptySilentMillis()) {
                skipCount++;
            }
        }
        if (subs.size() > skipCount) {
            return false;
        }
        for (Subscriber subscriber : subs) {
            subscriber.checkAndUpdatePushContext(task.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), task.getDataCenterPushCount());
        }
        PUSH_EMPTY_SKIP_COUNTER.inc();
        logger.info("[pushEmptySkip]{},{},{}", task.traceId, task.pushingTaskKey, task.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap());
        return true;
    }

    // some groupId not need to retry
    protected boolean needRetry(PushTask task, RetryReason reason) {
        return true;
    }

    private boolean retry(PushTask pushTask, RetryReason reason) {
        if (!needRetry(pushTask, reason)) {
            return false;
        }
        pushTask.retryCount++;
        final int retry = pushTask.retryCount;
        if (retry <= config.getPushTaskRetryTimes() && causeContinue(pushTask)) {
            final int backoffMillis = getRetryBackoffTime(retry);
            pushTask.expireAfter(backoffMillis);
            PUSH_RETRY_COUNTER.labels(reason.name()).inc();
            final boolean buffed = pushTaskWorkers.add(pushTask);
            logger.info("[retry]{},{},{},retry={},buffed={}", pushTask.traceId, pushTask.pushingTaskKey, pushTask.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), pushTask.retryCount, buffed);
            return buffed;
        }
        return false;
    }

    // check push empty, some group maybe could not tolerate push empty
    protected boolean interruptOnPushEmpty(DataCenterToSubDatumMap datum, PushData pushData, Cause cause, Subscriber sub, InetSocketAddress addr) {
        return false;
    }

    boolean doPush(PushTask pushTask) {
        if (!pushSwitchService.canIpPushMulti(pushTask.pushingTaskKey.inetSocketAddress.getAddress().getHostAddress(), pushTask.dataCenterToSubDatumMap.dataCenterSet())) {
            return false;
        }

        try {
            pushTask.trace.startPush();
            if (!checkPushRunning(pushTask)) {
                return false;
            }

            if (!causeContinue(pushTask)) {
                return false;
            }

            final PushData pushData = pushTask.createPushData();
            pushTask.setDataCenterPushInfos(pushData.getDataCenterPushInfo());

            if (interruptOnPushEmpty(pushTask.dataCenterToSubDatumMap, pushData, pushTask.trace.cause, pushTask.subscriber, pushTask.pushingTaskKey.inetSocketAddress)) {
                return false;
            }

            // double check
            if (!causeContinue(pushTask)) {
                return false;
            }
            // check push empty can skip (last push is also empty)
            if (checkSkipPushEmptyAndUpdateVersion(pushTask)) {
                return false;
            }

            pushingPushingTaskKeyToRecordMap.put(pushTask.pushingTaskKey, new PushRecord(pushTask.trace, pushTask.traceId, pushTask.retryCount, pushData.getDataCenterPushInfo()));
            clientNodeService.pushWithCallback(pushData.getPayload(), pushTask.subscriber.getClientUrl(), new PushClientCallback(pushTask));
            PUSH_CLIENT_ING_COUNTER.inc();
            logger.info("[pushing]{},{},{},{}", pushTask.traceId, pushTask.pushingTaskKey, pushTask.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), pushTask.getPushDataCount());
            return true;
        } catch (Throwable e) {
            handleDoPushException(pushTask, e);
        }
        return false;
    }

    void handleDoPushException(PushTask pushTask, Throwable e) {
        // try to delete self
        pushingPushingTaskKeyToRecordMap.remove(pushTask.pushingTaskKey);
        if (e instanceof RequestChannelClosedException) {
            pushTask.trace.finishPush(Trace.PushStatus.ChanClosed, pushTask.traceId, pushTask.getDataCenterPushInfos(), pushTask.retryCount);
            // channel closed, just warn
            logger.warn("[PushChanClosed]taskId={}, {}, {}", pushTask.traceId, pushTask.pushingTaskKey, e.getMessage());
            return;
        }

        // record push exception
        if (circuitBreakerRecordWhenDoPushError(pushTask.dataCenterToSubDatumMap, pushTask.subscriber.getClientUrl().getIp())) {
            for (Subscriber subscriber : pushTask.registerIdToSubscriberMap.values()) {
                if (!circuitBreakerService.onPushFail(pushTask.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), subscriber)) {
                    logger.info("[handleDoPushException]taskId={}, {}", pushTask.traceId, pushTask.pushingTaskKey);
                }
            }
        }
        if (e instanceof ChannelOverflowException) {
            retry(pushTask, RetryReason.Overflow);
            pushTask.trace.finishPush(Trace.PushStatus.ChanOverflow, pushTask.traceId, pushTask.getDataCenterPushInfos(), pushTask.retryCount);
            logger.error("[PushChanOverflow]taskId={}, {}, {}", pushTask.traceId, pushTask.pushingTaskKey, e.getMessage());
            return;
        }
        pushTask.trace.finishPush(Trace.PushStatus.Fail, pushTask.traceId, pushTask.getDataCenterPushInfos(), pushTask.retryCount);
        logger.error("[PushFail]taskId={}, {}", pushTask.traceId, pushTask.pushingTaskKey, e);
    }

    boolean circuitBreakerRecordWhenDoPushError(DataCenterToSubDatumMap datum, String ip) {
        return false;
    }

    int getRetryBackoffTime(int retry) {
        return BackOffTimes.getBackOffMillis(retry, config.getPushDataTaskRetryFirstDelayMillis(), config.getPushDataTaskRetryIncrementDelayMillis());
    }

    protected enum RetryReason {
        Waiting, Error, Overflow,
    }

    private static final class PushRecord {
        final Trace trace;
        final TraceId traceId;
        final int retryCount;

        private Map<String, DataCenterPushInfo> stringDataCenterPushInfoMap;

        PushRecord(Trace trace, TraceId traceId, int retryCount, Map<String, DataCenterPushInfo> stringDataCenterPushInfoMap) {
            this.trace = trace;
            this.traceId = traceId;
            this.retryCount = retryCount;
            this.stringDataCenterPushInfoMap = stringDataCenterPushInfoMap;
        }
    }

    final class Cleaner extends LoopExecuteTask {
        @Override
        public void _execute_() {
            int cleans = cleanPushingTaskRunTooLong();
            logger.info("cleans={}, callbackDiscardCounter={}, buffer={}, pushing={}", cleans, rejectedDiscardHandler.getDiscardCountThenReset(), pushTaskWorkers.size(), pushingPushingTaskKeyToRecordMap.size());
        }

        @Override
        public void _wait_() {
            ConcurrentUtils.sleepUninterruptibly(1000, TimeUnit.MILLISECONDS);
        }
    }

    class PushTaskImpl extends PushTask implements Runnable {
        PushTaskImpl(Cause cause, InetSocketAddress inetSocketAddress, Map<String, Subscriber> registerIdToSubscriberMap, DataCenterToSubDatumMap dataCenterToSubDatumMap) {
            super(cause, inetSocketAddress, registerIdToSubscriberMap, dataCenterToSubDatumMap);
        }

        protected PushData createPushData() {
            return pushDataGenerator.createPushData(DatumUtils.decompressDataCenterToSubDatumMap(dataCenterToSubDatumMap), registerIdToSubscriberMap);
        }

        @Override
        protected boolean commit() {
            try {
                if (!pushSwitchService.canIpPushMulti(pushingTaskKey.inetSocketAddress.getAddress().getHostAddress(), dataCenterToSubDatumMap.dataCenterSet())) {
                    return false;
                }
                // keyed by client.addr && (pushingKey%concurrencyLevel)
                // avoid generating too many pushes for the same client at the same time
                final int level = config.getClientNodePushConcurrencyLevel();
                pushExecutor.execute(new Tuple(pushingTaskKey.inetSocketAddress, pushingTaskKey.hashCode() % level), this);
                COMMIT_COUNTER.inc();
                return true;
            } catch (Throwable e) {
                logger.error("failed to exec push task {},{}", traceId, pushingTaskKey, e);
                return false;
            }
        }

        @Override
        public void run() {
            doPush(this);
        }
    }

    final class PushClientCallback implements Callback {
        final PushTask pushTask;

        PushClientCallback(PushTask pushTask) {
            this.pushTask = pushTask;
        }

        @Override
        public void onSuccess(Channel channel, Object message) {
            pushingPushingTaskKeyToRecordMap.remove(pushTask.pushingTaskKey);
            for (Subscriber subscriber : pushTask.registerIdToSubscriberMap.values()) {
                if (!circuitBreakerService.onPushSuccess(pushTask.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), pushTask.getDataCenterPushCount(), subscriber)) {
                    logger.info("PushY, but failed to updateVersion, {}, {}", pushTask.traceId, pushTask.pushingTaskKey);
                }
            }
            this.pushTask.trace.finishPush(Trace.PushStatus.OK, pushTask.traceId, pushTask.getDataCenterPushInfos(), pushTask.retryCount);
        }

        @Override
        public void onFailed(Channel channel, Throwable exception) {
            pushingPushingTaskKeyToRecordMap.remove(pushTask.pushingTaskKey);

            boolean needRecord = true;
            final boolean channelConnected = channel.isConnected();
            if (channelConnected) {
                retry(pushTask, RetryReason.Error);
            }

            if (exception instanceof InvokeTimeoutException) {
                this.pushTask.trace.finishPush(Trace.PushStatus.Timeout, pushTask.traceId, pushTask.getDataCenterPushInfos(), pushTask.retryCount);
                logger.error("[PushTimeout]taskId={}, {}", pushTask.traceId, pushTask.pushingTaskKey);
            } else {
                if (channelConnected) {
                    this.pushTask.trace.finishPush(Trace.PushStatus.Fail, pushTask.traceId, pushTask.getDataCenterPushInfos(), pushTask.retryCount);
                    logger.error("[PushFailed]taskId={}, {}", pushTask.traceId, pushTask.pushingTaskKey, exception);
                } else {
                    needRecord = false;
                    this.pushTask.trace.finishPush(Trace.PushStatus.ChanClosed, pushTask.traceId, pushTask.getDataCenterPushInfos(), pushTask.retryCount);
                    // channel closed, just warn
                    logger.warn("[PushChanClosed]taskId={}, {}", pushTask.traceId, pushTask.pushingTaskKey);
                }
            }

            if (needRecord) {
                // record push fail
                for (Subscriber subscriber : pushTask.registerIdToSubscriberMap.values()) {
                    if (!circuitBreakerService.onPushFail(pushTask.dataCenterToSubDatumMap.getDataCenterToSubDatumVersionMap(), subscriber)) {
                        logger.info("PushN, failed to do onPushFail, {}, {}", pushTask.traceId, pushTask.pushingTaskKey);
                    }
                }
            }
        }

        @Override
        public Executor getExecutor() {
            return pushCallbackExecutor;
        }
    }
}
